from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any, override

from tqdm import tqdm

from .interface import Processor
from .tokenizer import TokenizedDocumentStage

__all__ = [
    "LDATopic",
    "LDAPass",
    "GensimLDAVisualizationPreset",
    "LDAStatistics",
    "GensimLDAStage",
    "GensimLDAProcessor",
]

type LDATopic = Sequence[tuple[str, float]]
type LDAPass = Sequence[LDATopic]


@dataclass(frozen=True)
class GensimLDAVisualizationPreset(object):
    num_topics: int
    visualization_data: Any


@dataclass(frozen=True)
class LDAStatistics(object):
    perplexity: float
    coherence: float


@dataclass(frozen=True)
class GensimLDAStage(object):
    visualization_presets: Sequence[GensimLDAVisualizationPreset]
    passes: Sequence[LDAPass]
    statistics: Sequence[LDAStatistics]


class GensimLDAProcessor(Processor[TokenizedDocumentStage, GensimLDAStage]):
    """
    A processor that performs LDA analysis on tokenized documents.

    This processor uses `gensim` to perform LDA analysis.
    """

    _config: dict

    def __init__(self, global_config: dict) -> None:
        self._config = global_config["topic"]["lda"]

    @override
    def process(self, stage: TokenizedDocumentStage) -> GensimLDAStage:
        # These imports cannot be placed at the top of the file, because they will trigger a deprecation warning.
        #
        # I don't know why this is happening, but we better do so.
        import pyLDAvis.gensim_models
        from gensim import corpora
        from gensim.models import CoherenceModel, LdaModel

        # prepare dictionary and corpus
        dictionary = corpora.Dictionary(stage.tokenized_documents)
        corpus = [dictionary.doc2bow(doc) for doc in stage.tokenized_documents]

        # LDA and coherence analysis
        topic_nums = range(self._config["min_topics"], self._config["max_topics"] + 1)
        visualization_presets = []
        passes = []
        staticstics = []
        for num_topics in tqdm(topic_nums, desc="LDA Analysis"):
            # calculate perplexity score
            lda = LdaModel(
                corpus,
                num_topics=num_topics,
                id2word=dictionary,
                passes=self._config["passes"],
                random_state=self._config["random_state"],
            )
            perplexity = lda.log_perplexity(corpus)

            # calculate coherence score
            coherence_model = CoherenceModel(
                model=lda,
                texts=stage.tokenized_documents,
                dictionary=dictionary,
            )
            coherence = coherence_model.get_coherence()

            # store results
            visualization_presets.append(
                GensimLDAVisualizationPreset(
                    num_topics,
                    pyLDAvis.gensim_models.prepare(lda, corpus, dictionary, mds="mmds"),
                )
            )
            current_pass = []
            for _, pair in lda.show_topics(num_topics=num_topics, formatted=False):
                current_pass.append(pair)
            passes.append(current_pass)
            staticstics.append(LDAStatistics(perplexity, coherence))
        return GensimLDAStage(visualization_presets, passes, staticstics)
